"""
蚁群算法 TSP
"""
import os

os.chdir(os.path.dirname(__file__))

import tsplib95
import random
from tsplib95.models import Problem
import numpy as np
import matplotlib.pyplot as plt #导入所需要的库


class ACO(object):
    def __init__(self, generations: int, alpha: float, beta: float, rho: float, Q: int, ant_count=-1, strategy=0):
        """
        :param generations: 最大迭代次数
        :param alpha: 信息启发因子，反映了蚁群在运动过程中残留信息量的相对重要程度
        :param beta: 期望启发因子，反映了期望值（距离）的相对重要程度
        :param rho: 信息素挥发系数
        :param Q: 常量，表示蚂蚁循环一个过程在经过的路径上锁匙放的信息素总量
        :param ant_count: 蚁群中蚂蚁的总数，不设置的话默认为问题规模的1.5倍？？？？？
        :param strategy: 信息素更新策略：0-蚁周模型（ant-cycle）默认, 1-蚁量模型（ant-quality）, 2-蚁密模型（ant-density）
        """
        self.Q = Q
        self.rho = rho
        self.beta = beta
        self.alpha = alpha
        self.ant_count = ant_count
        self.generations = generations
        self.update_strategy = strategy
        self.pheromone = [[]]   # 路径信息素含量，传入问题时初始化
        self.n = 0              # 问题规模，传入问题时初始化

    def _update_pheromone(self, ants: list):
        """
        迭代一轮之后，更新各边信息素含量：挥发一部分，
        :param ants: 勤劳的蚂蚁们
        """
        for i, row in enumerate(self.pheromone):
            for j, col in enumerate(row):
                self.pheromone[i][j] *= (1-self.rho)  # 挥发后残留的信息素
                for ant in ants:
                    # 蚂蚁走过后的信息素增量
                    self.pheromone[i][j] += ant.pheromone_delta[i][j]

    def solve(self, problem: Problem):
        """
        进行问题求解
        :param problem: tsplib95中定义的问题模型
        求解返回的路径是节点位置索引，不是节点编号
        """

        # 问题规模
        self.n = problem.dimension
        # 索引和节点的对应关系
        self.nodes = list(problem.get_nodes())

        # 蚁群规模
        if(self.ant_count < 0):
            self.ant_count = int(1.5 * self.n)
        # 初始化信息素
        self.pheromone = [[1 / (self.n * self.n)
                           for j in range(self.n)] for i in range(self.n)]
        # 当前最优解（最小值）
        best_cost = float('inf')
        # 最优解路径
        best_solution = []
        cost_log = []
        best_cost_log = []
        # 终止条件：迭代到最大代数
        for gen in range(self.generations):
            # 初始化蚁群
            ants = [_Ant(self, problem) for i in range(self.ant_count)]
            gen_best_cost = float('inf')
            gen_best_solution = []
            # 此处按蚂蚁迭代，可以拆分成多线程形式
            # 如果按每只蚂蚁走一步，轮流来则不太好实现多线程
            for ant in ants:
                for i in range(self.n - 1):  # 初始化占用1个节点，因此只需再选择n-1个节点
                    ant._select_next()
                # 从最后一个节点，回到起点
                ant.total_cost += problem.get_weight(self.nodes[ant.tabu[-1]], self.nodes[ant.tabu[0]])
                if ant.total_cost < gen_best_cost:
                    gen_best_cost = ant.total_cost
                    gen_best_solution = [] + ant.tabu
                # 蚂蚁走完后，计算其留下的的信息素
                ant._update_pheromone_delta()

            # 一次迭代之后，更新全局的信息素
            self._update_pheromone(ants)
            if gen_best_cost < best_cost:
                best_cost = gen_best_cost
                best_solution = gen_best_solution
            # 输出，可以用这个做可视化
            
            cost_log.append(gen_best_cost)
            best_cost_log.append(best_cost)
            print(f"gen {gen} best_cost {best_cost} gen_best_cost {gen_best_cost}")

        return best_solution, best_cost, cost_log, best_cost_log


class _Ant(object):
    def __init__(self, aco: ACO, problem: Problem):
        """
        :param aco: 蚁群
        :param problem: tsplib95中定义的问题模型
        """
        self.problem = problem
        self.colony = aco
        self.total_cost = 0.0  # 代价
        self.tabu = []  # 禁忌表，记录本次迭代中访问过的节点
        self.pheromone_delta = []  # 当前蚂蚁在本次循环中留在路径上信息素的量
        # 初始化允许访问的节点
        self.allowed = [i for i in range(aco.n)]
        self.eta = [[0 if i == j else 1 / (self.problem.get_weight(i+1, j+1))
                     for j in range(aco.n)] for i in range(aco.n)]  # 期望值
        start = random.randint(0, aco.n - 1)  # 随机设置蚂蚁的初始位置
        self.tabu.append(start)
        self.current = start
        self.allowed.remove(start)

    def _select_next(self):
        denominator = 0
        for i in self.allowed:
            denominator += self.colony.pheromone[self.current][i] ** self.colony.alpha * self.eta[self.current][i] ** self.colony.beta
        # 生成概率数组
        probabilities = [0 for i in range(self.colony.n)]
        for i in range(self.colony.n):
            try:
                self.allowed.index(i)
                probabilities[i] = self.colony.pheromone[self.current][i] ** self.colony.alpha * self.eta[self.current][i] ** self.colony.beta / denominator
            except ValueError:
                pass
        # 采用轮盘赌选择下一个节点
        selected = 0
        rand = random.random()
        for i, probability in enumerate(probabilities):
            rand -= probability
            if rand <= 0:
                selected = i
                break
        self.allowed.remove(selected)
        self.tabu.append(selected)
        self.total_cost += self.problem.get_weight(self.colony.nodes[self.current], self.colony.nodes[selected])
        self.current = selected

    def _update_pheromone_delta(self):
        """
        信息素更新策略
        """
        self.pheromone_delta = [
            [0 for j in range(self.colony.n)] for i in range(self.colony.n)]
        for _ in range(1, len(self.tabu)):
            i = self.tabu[_ - 1]
            j = self.tabu[_]
            # 判断更新策略
            if self.colony.update_strategy == 1:  # ant-quality system
                self.pheromone_delta[i][j] = self.colony.Q
            elif self.colony.update_strategy == 2:  # ant-density system
                self.pheromone_delta[i][j] = self.colony.Q / self.problem.get_weight(self.colony.nodes[i],self.colony.nodes[j])
            else:  # ant-cycle system
                self.pheromone_delta[i][j] = self.colony.Q / self.total_cost


# TSP问题数据
problem = tsplib95.load('../test_data/tsp/opt_tour/berlin52.tsp')



aco = ACO(20, 1.0, 4.0, 0.7, 150,300)
best_solution, best_cost, cost_log, best_cost_log = aco.solve(problem)
# 适应值函数曲线
plt.figure(1)
gen = [str(i+1) for i in range(len(cost_log))]
plt.plot(gen, cost_log, c='green', linestyle='--', label="generation best cost")
plt.plot(gen, best_cost_log, c='red', label="global best cost")
plt.legend(loc='best')


x = []
y = []
for point in problem.get_nodes():
    x.append(problem.get_display(point)[0])
    y.append(problem.get_display(point)[1])
plt.figure(2)
for _ in range(1, len(best_solution)):
    i = best_solution[_ - 1]
    j = best_solution[_]
    # noinspection PyUnresolvedReferences
    plt.quiver(x[i], y[i], x[j] - x[i], y[j] - y[i], color='r', width=0.005, angles='xy', scale=1,
                scale_units='xy')
plt.quiver(x[best_solution[-1]], y[best_solution[-1]], x[best_solution[0]] - x[best_solution[-1]], y[best_solution[0]] - y[best_solution[-1]], color='r', width=0.005, angles='xy', scale=1,
                scale_units='xy')
plt.show()